import pickle
import numpy as np
import torch
import os

from Config.Config import ACTNET200V13_PKL,TRAIN_PROPOSALS_PKL,TEST_PROPOSALS_PKL,VAL_PROPOSALS_PKL,TSNscore_FEATURE_PATH
from DataSet.Tools_TSNscore import load_TSNscore_from_file,extract_TSN_feature
from DataSet.Utils import mulity_scale_pairs

class Dataset_TSNscore_Pair(object):
    '''
    取数据, 分为test和train两种模式
    '''
    def __init__(self,subset,modality):
        '''
        :param modality:  test or train
        :param subset: testing validation traning
        '''
        self.modality = modality
        self.subset = subset

        self.load_groundtruth()
        self.load_proposals()

    def load_groundtruth(self):
        with open(ACTNET200V13_PKL,'rb') as f:
            groundtruth = pickle.load(f)['database']
        # filter by subset
        remainkeyset = []
        for key in groundtruth.keys():
            if groundtruth[key]['subset'] == self.subset:
                remainkeyset.append(key)
        self.groundtruth = dict()
        for key in remainkeyset:
            self.groundtruth[key] = groundtruth[key]
        self.vids = remainkeyset

    def load_proposals(self):
        PKL = None
        if self.subset == 'training':
            PKL = TRAIN_PROPOSALS_PKL
        elif self.subset == 'testing':
            PKL = TEST_PROPOSALS_PKL
        elif self.subset == 'validation':
            PKL = VAL_PROPOSALS_PKL
        with open(PKL,'rb') as f:
            self.proposals = pickle.load(f)


    def _random_sample_traing_item(self,batch=2048, random=512, hard_neg=512, maxnump=32, PYRAMID=[12,12], thr_ious=[0.7,0.3]):
        '''
        无穷无尽的取训练用的batch
        :param batch:
        :param random:
        :param hard_neg:
        :param maxnump:
        :param PYRAMID:
        :param thr_ious:
        :return:
        '''
        target_pos = batch-hard_neg-random
        pos_count = 0
        neg_count = 0
        rand_count = 0
        count = 0

        x = np.zeros((batch, 1, np.sum(PYRAMID), 202))
        y = np.ones(batch)*-1

        for pos,vid in enumerate(self.vids):

            vitem = self.groundtruth[vid]
            if not os.path.exists(TSNscore_FEATURE_PATH+'/feat/{}.h5'.format(vid)) or vid not in self.proposals or len(self.proposals[vid]) == 0:
                continue
            feat = load_TSNscore_from_file(vid)
            num_frame = vitem['numf']
            duration = vitem['duration']
            proposal = np.asarray(self.proposals[vid])

            max_ious = proposal[:,:,2].max(axis=1)
            max_ious_lv = proposal[:,:,2].argmax(axis=1)

            pos_indexes = np.nonzero(max_ious>thr_ious[0])[0]
            neg_indexes = np.nonzero(max_ious<thr_ious[1])[0]

            pos_pairs = []
            for i in range(len(pos_indexes)):
                ind = pos_indexes[i]
                lv = max_ious_lv[ind]
                pos_pairs += mulity_scale_pairs(proposal[ind,], [lv])

            np.random.shuffle(pos_pairs)
            for pair in pos_pairs[:min(target_pos,maxnump)]:
                if pos_count < target_pos and count < batch:
                    data = extract_TSN_feature(feat, pair, num_frame, duration, PYRAMID=PYRAMID)
                    if data is not None and data.shape[0]==np.sum(PYRAMID):
                        x[count, ...].flat = data.flat
                        y[count, ...] = 1
                        pos_count += 1
                        count +=1

            # sample random negative around around positive
            rand_pairs = []
            for i in range(len(pos_indexes)):
                ind = pos_indexes[i]
                lv = max_ious_lv[ind]
                rand_pairs += mulity_scale_pairs(proposal[ind,], [i for i in range(6) if i!=lv])
            np.random.shuffle(rand_pairs)

            for pair in rand_pairs[:min(random,maxnump)]:
                if rand_count<random and count < batch:
                    data = extract_TSN_feature(feat, pair, num_frame, duration, PYRAMID=PYRAMID)
                    if data is not None and data.shape[0] == np.sum(PYRAMID):
                        x[count, ...].flat = data.flat
                        y[count, ...] = 0
                        rand_count += 1
                        count += 1

            # sample hard negatives
            neg_pairs = []
            for i in range(len(neg_indexes)):
                ind = neg_indexes[i]
                neg_pairs += mulity_scale_pairs(proposal[ind,])
            np.random.shuffle(neg_pairs)

            for pair in neg_pairs[:min(hard_neg,maxnump)]:
                if neg_count < hard_neg and count < batch:
                    data = extract_TSN_feature(feat, pair, num_frame, duration, PYRAMID=PYRAMID)
                    if data is not None and data.shape[0]==np.sum(PYRAMID):
                        x[count, ...].flat = data.flat
                        y[count, ...] = 0
                        neg_count += 1
                        count += 1

            if count == batch:
                # 位移动一下
                self.vids = self.vids[pos:] + self.vids[:pos]
                break

        return x,y

    def _enum_every_test_item(self):
        # 一次出来一组6个和对应的VID
        for vid,proposals in self.proposals.items():

            vitem = self.groundtruth[vid]
            if not os.path.exists(TSNscore_FEATURE_PATH+'/feat/{}.h5'.format(vid)) or vid not in self.proposals or len(self.proposals[vid]) == 0:
                continue
            feat = load_TSNscore_from_file(vid)
            num_frame = vitem['numf']
            duration = vitem['duration']
            proposals = np.asarray(proposals)

            # 64 * 6
            for i in range(proposals.shape[0]):
                current = np.copy(proposals[i,...])
                # current shape 6x4
                current[:,2] = 0
                pairs = mulity_scale_pairs(current)
                ret = []
                index = []
                for pos,pair in enumerate(pairs):
                    x = extract_TSN_feature(feat,pair,num_frame,duration,PYRAMID=[12,12])
                    if x is not None and x.shape[0] == 24:
                        # x1,x2 = x[:12,:],x[12:,:]
                        ret.append(x)
                        index.append(pos)
                # TODO 兼容有问题
                # 由于特征提取的原因, 有时候并不能正好得到6张图片
                ret = np.asarray(ret)
                current = current[index,...]
                yield [vid,current,ret]

    def nextbatch(self):
        return self._random_sample_traing_item()

    def enum_genfucntion(self):
        return self._enum_every_test_item()

def test():
    dataset = Dataset_TSNscore_Pair('validation', 'train')
    x,y = dataset.nextbatch()

if __name__=='__main__':
    dataset = Dataset_TSNscore_Pair('validation', 'test')
    gen = dataset.enum_genfucntion()
    # b = next(gen)
    # b[2].shape
    # cnt = 0
    # for b in next(gen):
    #     cnt += 1
    #     if cnt > 100:
    #         break

